import cv2
import numpy as np


class DocScanner:
    def __init__(self):
        pass

    def load_image(self, file_path):
        # 读取图像，转换为灰度图像
        image = cv2.imread(file_path)
        if image is None:
            raise FileNotFoundError("无法读取文件: " + file_path)
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 提高对比度
        gray = cv2.equalizeHist(gray)

        # 使用Canny边缘检测
        edged = cv2.Canny(gray, 50, 200, apertureSize=3)

        # 寻找轮廓
        contours, _ = cv2.findContours(edged.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

        contours = sorted(contours, key=cv2.contourArea, reverse=True)[:5]

        # 寻找文档的四个角落
        for c in contours:
            p = cv2.arcLength(c, True)
            approx = cv2.approxPolyDP(c, 0.02 * p, True)

            if len(approx) == 4:
                return image, approx.reshape(4, 2)

        # 如果找不到四个角落，返回原始图像和None
        return image, None

    def order_points(self, pts):
        """对给定的四个点进行排序，返回排序后的点。"""
        rect = np.zeros((4, 2), dtype="float32")
        s = pts.sum(axis=1)
        rect[0] = pts[np.argmin(s)]
        rect[2] = pts[np.argmax(s)]
        diff = np.diff(pts, axis=1)
        rect[1] = pts[np.argmin(diff)]
        rect[3] = pts[np.argmax(diff)]
        return rect

    def transform(self, img, corners):
        """根据四个角落裁剪图片。"""
        if corners is None:
            return img

        corners = self.order_points(corners)  # 新增的排序步骤

        # 计算目标图片的宽度和高度
        width, height = self.get_image_dimensions(corners)

        # 根据以上宽度和高度创建于原有四个点顺序对应的目标四个点的坐标
        dst_points = np.array([
            [0, 0],
            [width - 1, 0],
            [width - 1, height - 1],
            [0, height - 1]], dtype="float32")

        # 通过投影变换展平文档
        M = cv2.getPerspectiveTransform(corners, dst_points)  # 获取透视变换矩阵
        dst = cv2.warpPerspective(img, M, (width, height))  # 对图片进行透视变换
        return dst  # 返回展平后的图像

    def get_image_dimensions(self, corners):
        """计算图片的宽度和高度。"""
        top_left_corner, top_right_corner, bottom_right_corner, bottom_left_corner = corners

        # 计算宽度
        width_A = np.sqrt(((bottom_right_corner[0] - bottom_left_corner[0]) ** 2) +
                          ((bottom_right_corner[1] - bottom_left_corner[1]) ** 2))
        width_B = np.sqrt(((top_right_corner[0] - top_left_corner[0]) ** 2) +
                          ((top_right_corner[1] - top_left_corner[1]) ** 2))
        width = max(int(width_A), int(width_B))

        # 计算高度
        height_A = np.sqrt(((top_right_corner[0] - bottom_right_corner[0]) ** 2) +
                           ((top_right_corner[1] - bottom_right_corner[1]) ** 2))
        height_B = np.sqrt(((top_left_corner[0] - bottom_left_corner[0]) ** 2) +
                           ((top_left_corner[1] - bottom_left_corner[1]) ** 2))
        height = max(int(height_A), int(height_B))

        return width, height
